from sys import flags
from time import time
import os,shutil
from struct import pack,unpack,calcsize
import CreateModelMap

PresentBytes = ""
Presentlocation = 0


def LoadByteValue(fmt):
    global PresentBytes
    global Presentlocation
    r_str = unpack(fmt,PresentBytes[:calcsize(fmt)])
    PresentBytes = PresentBytes[calcsize(fmt):]
    Presentlocation += calcsize(fmt)
    if len(r_str) > 1:
        return r_str
    else:
        return r_str[0]

class MDXBlock:
    blockname = ""
    def __init__(self,name,model):
        self.origin = Presentlocation
        if LoadByteValue("4s").decode() == name:
            self.blockname = name
            self.model = model
    
    def ReadBounds(self):
        r_dist = {}
        r_dist['BoundsRadius'] = LoadByteValue('f')
        r_dist['MinExt'] = {
            'x' : LoadByteValue('f'),
            'y' : LoadByteValue('f'),
            'z' : LoadByteValue('f'),
        }
        r_dist['MaxExt']= {
            'x' : LoadByteValue('f'),
            'y' : LoadByteValue('f'),
            'z' : LoadByteValue('f'),
        }
        return r_dist
        
class GetModelValue:
    # 固有属性
    def __init__(self,filename):
        with open(filename,"rb") as fi:
            self.filename = filename
            self.Bytes = fi.read()

        global PresentBytes
        global Presentlocation
        PresentBytes = self.Bytes
        Presentlocation = 0

        self.Block = []

        if LoadByteValue("4s").decode() == "MDLX":
            self.type = 'model'
            self.VERS = VERS(self.Bytes,self)
            self.MODL = MODL(self.Bytes,self)
            self.SEQS = SEQS(self.Bytes,self)
            self.GLBS = GLBS(self.Bytes,self)
            self.MTLS = MTLS(self.Bytes,self)
            self.TEXS = TEXS(self.Bytes,self)
        
    def ExportModel(self,filename):
        # UpdateValue 
        # TEXS
        if not self.TEXS.is_None:
            if len(self.TEXS.change_tuple[0]) > 0:
                for i in range(len(self.TEXS.change_tuple[0])):
                    self.Bytes = self.Bytes.replace(self.TEXS.change_tuple[0][i],self.TEXS.change_tuple[1][i],1)
                self.TEXS.change_tuple = ([],[])

        with open(filename,"wb") as fi:
            fi.write(self.Bytes)

class VERS(MDXBlock):
    def __init__(self,Bytes,model):
        if b'VERS' != unpack('4s',PresentBytes[:4])[0]:
            return
        MDXBlock.__init__(self,"VERS",model)
        self.size = LoadByteValue('l')
        self.version = LoadByteValue('l')
        self.final = Presentlocation
        self.Bytes = Bytes[self.origin:self.final]

class MODL(MDXBlock):
    def __init__(self,Bytes,model):
        global PresentBytes
        global Presentlocation
        if b'MODL' != unpack('4s',PresentBytes[:4])[0]:
            return
        MDXBlock.__init__(self,"MODL",model)
        self.size = LoadByteValue('l')
        # self.name = LoadByteValue('80s')
        # self.animationFileName = LoadByteValue('260s').decode()
        # self.extent = MDXBlock.ReadBounds(self)
        # self.blendtime = LoadByteValue("f")

        PresentBytes = PresentBytes[self.size:]
        Presentlocation += self.size

        self.final = Presentlocation
        self.Bytes = Bytes[self.origin:self.final]

class SEQS(MDXBlock):
    def __init__(self,Bytes,model):
        global PresentBytes
        global Presentlocation
        if b'SEQS' != unpack('4s',PresentBytes[:4])[0]:
            return
        MDXBlock.__init__(self,"SEQS",model)
        self.size = LoadByteValue('l')

        # self.seqs = []
        # for _ in range(int(self.size / 132)):
        #     r_dist = {
        #         'name' : LoadByteValue('80s'),
        #         'interval' : {
        #             'start' : LoadByteValue('l'),
        #             'end'   : LoadByteValue('l'),
        #         },
        #         'moveSpeed' : LoadByteValue("f"),
        #         'is_looping' : LoadByteValue('l'), #0 looping
        #                                            #1 not looping
        #         'rarity' : LoadByteValue("f"),
        #         'syncPoint' : LoadByteValue('l'),
        #         'extent' : MDXBlock.ReadBounds(self),
        #     }
        #     self.seqs.append(r_dist)

        PresentBytes = PresentBytes[self.size:]
        Presentlocation += self.size

        self.final = Presentlocation
        self.Bytes = Bytes[self.origin:Presentlocation]

class GLBS(MDXBlock):
    def __init__(self,Bytes,model):
        global PresentBytes
        global Presentlocation
        if b'GLBS' != unpack('4s',PresentBytes[:4])[0]:
            return
        MDXBlock.__init__(self,"GLBS",model)
        self.size = LoadByteValue('l')
        # self.glb = []
        # for _ in range(int(self.size / 4)):
        #     self.glb.append(LoadByteValue('l'))

        PresentBytes = PresentBytes[self.size:]
        Presentlocation += self.size
        
        self.final = Presentlocation
        self.Bytes = Bytes[self.origin:self.final]


# MTLS_Subset
# 2021/7/10 skip
# class KMTF(MDXBlock):
#     def __init__(self,model):
#         if b'KMTF' != unpack('4s',PresentBytes[:4])[0]:
#             return
#         MDXBlock.__init__(self,'KMTF',model)
    
# class KMTA(MDXBlock):
#     def __init__(self,model):
#         if b'KMTA' != unpack('4s',PresentBytes[:4])[0]:
#             return
#         MDXBlock.__init__(self,'KMTA',model)

# class LAYS(MDXBlock):
#     def __init__(self,model):
#         if b'LAYS' != unpack('4s',PresentBytes[:4])[0]:
#             return
#         MDXBlock.__init__(self,'LAYS',model)
#         self.layersCount = LoadByteValue('l')
#         r_is_800version = model.VERS.version > 800
#         r_is_900version = model.VERS.version >= 900 #重制版 先不管了报错再说
#         self.layer = []
#         for _ in range(self.layersCount):
#             r_dist = {
#                 'inclusivesize' : LoadByteValue('l'),
#                 'filterMode'    : LoadByteValue('l'),
#                                         # 0 none
#                                         # 1 transparent
#                                         # 2 blend
#                                         # 3 additive
#                                         # 4 add alpha
#                                         # 5 modulate
#                                         # 6 modulate 2x
#                 'shadingFlags'   : LoadByteValue('l'),
#                 'textureId'       : LoadByteValue('l'),
#                 'textureAnimationId'  : LoadByteValue('l'),
#                 'coordId'         : LoadByteValue('l'),
#                 'alpha'           : LoadByteValue('f'),
#             }
        
#             if r_is_800version :
#                 r_dist['emissiveGain'] = LoadByteValue('f')
#                 r_dist['fresnelColor'] = {
#                     'R' : LoadByteValue('f'),
#                     'G' : LoadByteValue('f'),
#                     'B' : LoadByteValue('f'),
#                 }
#                 r_dist['fresnelOpacity'] = LoadByteValue('f')
#                 r_dist['fresnelTeamColor'] = LoadByteValue('f')
            
#             self.KMTF = KMTF(model)
#             self.KMTA = KMTA(model)
            
#             if r_is_800version :
#                 if b'KMTE' == unpack('4s',PresentBytes[:4])[0]:
#                     print('检测到KMTE标记')
#                     return
            
#             if r_is_900version :
#                 if b'KFC3' == unpack('4s',PresentBytes[:4])[0]:
#                     print('检测到KMTE标记')
#                     return
#                 elif b'KFCA' == unpack('4s',PresentBytes[:4])[0]:
#                     print('检测到KFCA标记')
#                     return
#                 elif b'KFTC' == unpack('4s',PresentBytes[:4])[0]:
#                     print('检测到KFTC标记')
#                     return
            
#             self.layer.append(r_dist)

# ————————————————————————————————————————————————

class MTLS(MDXBlock):
    def __init__(self,Bytes,model):
        global PresentBytes
        global Presentlocation
        if b'MTLS' != unpack('4s',PresentBytes[:4])[0]:
            return
        MDXBlock.__init__(self,'MTLS',model)
        self.size = LoadByteValue('l')
        
        #直接TM的进行一个跳的过 (将来回来再补坑了
        # r_is_800version = model.VERS.version > 800
        # self.Material = []
        
        # while self.size > Presentlocation - (self.origin + 8):
        #     r_dist = {
        #         'inclusivesize' : LoadByteValue('l'),
        #         'priorityplane' : LoadByteValue('l'),
        #         'flags' : LoadByteValue('l'),
        #     }
        #     if r_is_800version:
        #         r_dist['shader'] = LoadByteValue('80s')

        #     r_dist['LAYS'] = LAYS(model)
        #     self.Material.append(r_dist)
        PresentBytes = PresentBytes[self.size:]
        Presentlocation += self.size
        
        self.final = Presentlocation
        self.Bytes = Bytes[self.origin:self.final]

class TEXS(MDXBlock):
    is_None = True
    def __init__(self,Bytes,model):
        if b'TEXS' != unpack('4s',PresentBytes[:4])[0]:
            return
        MDXBlock.__init__(self,'TEXS',model)
        self.is_None = False
        self.change_tuple = ([],[])
        self.size = LoadByteValue('l')
        self.texture = []
        for _ in range(int(self.size / 268)):
            r_dist = {
                'replaceableId' : LoadByteValue('l'),
                'filename' : LoadByteValue('260s'),
                'flags' : LoadByteValue('l'),
            }
            self.texture.append(r_dist)
        self.final = Presentlocation
        self.Bytes = Bytes[self.origin:self.final]
        pass

    def GetTexturePath(self):
        r_list = []
        for i in range(int(self.size / 268)):
            r_list.append(self.texture[i]['filename'].decode('gbk').rstrip('\x00'))

        return r_list

    def ReviseTexturePath(self,number,texturepath):
        old_str = self.texture[number]['filename']
        new_str = pack('260s',texturepath.encode("utf-8"))
        self.change_tuple[0].append(old_str)
        self.change_tuple[1].append(new_str)

# 在返回文件名和获取mdx贴图路径时候做全小写处理
def is_sam_path(modelpath):
    r_tuple = ([],[])
    if not os.path.isdir(modelpath):
        modelpath = "."
    for root, dirs, files in os.walk(modelpath):
        for file in files:
            r_tuple[0].append(file) ,r_tuple[1].append((root + "\\" + file))
    return r_tuple

def GetFileName(path):
    number = path.rfind("\\")
    if number == -1:
        return path
    else:
        return path[number + 1:]

def GetFilePath(path):
    number = path.rfind("\\")
    if number == -1:
        return path
    else:
        return path[:number]

def GetFiletip(path):
    number = path.find("\\")
    if number == -1:
        return path
    else:
        return path[:number]

def is_native_texture(path):
    return GetFiletip(path).lower() in ['abilities','buildings','doodads','environment','font','fonts','maps','objects','pathtextures','replaceabletextures','scripts','sharedmodels','sound','splats','terrainart','textures','units','ui']

# lower
def is_str_inlist(r_str,r_list):
    for i in range(len(r_list)):
        r_list[i] = r_list[i].lower()
    return r_str.lower() in r_list

# lower
def GetSerListIndex(r_str,r_list):
    for i in range(len(r_list)):
        r_list[i] = r_list[i].lower()
    return r_list.index(r_str.lower())

# War3ImageSizeFilter
def is_error_blp(filePath):
    if filePath[-4:] == '.blp':
        with open(filePath,'rb') as fi:
            Bytes = fi.read()
        
        ImagesSize = unpack('2i',Bytes[12:20])

        return not (bin(ImagesSize[0]).count("1") == 1 and bin(ImagesSize[1]).count("1") == 1)
    return False

def is_error_tga(filePath):
    if filePath[-4:] == '.tga':
        with open(filePath,'rb') as fi:
            Bytes = fi.read()
        
        ImagesSize = unpack('2h',Bytes[12:16])

        return not (bin(ImagesSize[0]).count("1") == 1 and bin(ImagesSize[1]).count("1") == 1)
    return False

def is_error_model(textures,dirs_paths):
    for texture in textures:
        if GetFileName(texture) != "" and not is_str_inlist(GetFileName(texture),dirs_paths[0]) and not is_native_texture(texture):
            print(texture)
            return True,"NoneTexture"
        elif GetFileName(texture) != "" and not is_native_texture(texture) and is_error_blp(dirs_paths[1][GetSerListIndex(GetFileName(texture),dirs_paths[0])]):
            print(texture)
            return True,"BlpSizeError"
    return False

# 强打包(设置贴图位置)
def ViolenceReviseTexturePath(model,dirs_paths,textures):
    global textureNumber
    for texture in textures:
        if GetFileName(texture) != "" and is_str_inlist(GetFileName(texture),dirs_paths[0]):
            # 设置模型贴图路径
            r_texture =  str(textureNumber).rjust(6,"0") + " " + str(time_number) + GetFileName(texture)[-4:]
            r_path = dirs_paths[1][GetSerListIndex(GetFileName(texture),dirs_paths[0])]
            model.TEXS.ReviseTexturePath(textures.index(texture),r_texture)
            # 处理贴图命名与路径
            shutil.copyfile(r_path,MapPath + '\\lop\\resource\\' + r_texture)
            change_Table[0].append(r_texture)
            change_Table[1].append(r_path)
            textureNumber += 1

# 弱打包(复制文件)
def ViolenceReviseTexturePath(model,dirs_paths,textures):
    global textureNumber
    for texture in textures:
        if GetFileName(texture) != "" and is_str_inlist(GetFileName(texture),dirs_paths[0]):
            # 设置模型贴图路径
            r_texture = GetFileName(texture)
            r_path = dirs_paths[1][GetSerListIndex(GetFileName(texture),dirs_paths[0])]
            model.TEXS.ReviseTexturePath(textures.index(texture),r_texture)
            # 处理贴图命名与路径
            shutil.copyfile(r_path,MapPath + '\\lop\\resource\\' + r_texture)
            change_Table[0].append(r_texture)
            change_Table[1].append(r_path)
            textureNumber += 1

def UnpackMap(MapPath,MapName,is_ModelGroup,ModelGroupNumber):
    # 跳过
    if ModelGroupNumber in skiplist:
        shutil.rmtree(MapPath + '\\lop\\resource')
        os.mkdir(MapPath + '\\lop\\resource',777)
        return
    
    # 添加前缀
    if is_ModelGroup:
        MapName = MapName + "[%s]" % (str(ModelGroupNumber))

    # 输出日志
    with open(MapPath + '\\' + MapName + " change_Table.txt","w",encoding='utf-8') as fi:
        r_str = ""
        for i in range(len(change_Table[2])):
            r_str += change_Table[2][i] + "\t" + change_Table[3][i] + "\n"
        
        for i in range(len(change_Table[0])):
            r_str += change_Table[0][i] + "\t" + change_Table[1][i] + "\n"
            
        r_str += "ErrorModel:\n"
        for model in error_model:
            r_str += model[0] + "\t" + model[1] + "\n"
        fi.write(r_str)

    CreateModelMap.main(MapName + '.w3x')
    shutil.rmtree(MapPath + '\\lop\\resource')
    os.mkdir(MapPath + '\\lop\\resource',777)

work_path = ''
MapPath = ''
mdx_files = []
mdl_files = []
textureNumber = 0
time_number = int(time())
change_Table = ([],[],[],[])
error_model = []

def main(p_work_path,p_map_path,is_ModelGroup):
    global PresentBytes
    global Presentlocation
    global work_path
    global MapPath
    global textureNumber
    global time_number
    global mdx_files
    global mdl_files
    global change_Table
    global error_model

    work_path = p_work_path
    MapPath = p_map_path
    PresentBytes = ""
    Presentlocation = 0
    mdx_files = []
    mdl_files = []
    texture_files = {
        'name' : [],
        'path' : [],
    }
    error_model = []
    textureNumber = 0
    time_number = int(time())
    change_Table = ([],[],[],[])
    #  private 
    MapName = GetFileName(work_path)
    modelNumber = 1
    ModelGroupNumber = 1

    # init
    os.chdir(work_path)
    shutil.rmtree(MapPath + '\\lop\\resource')
    os.mkdir(MapPath + '\\lop\\resource',777)

    # 获取模型目录
    for root, dirs, files in os.walk("."):
        for file in files:
            if '.mdx' == file[-4:]:
                mdx_files.append((root + "\\" + file)[2:])
            elif '.mdl' == file[-4:]:
                mdl_files.append((root + "\\" + file)[2:])
            elif file[-4:] in ['.tga','.blp']:
                texture_files['name'].append(file)
                texture_files['path'].append((root + "\\" + file)[2:])

    iterantTexturePath = []
    for number in range(len(texture_files['name'])):
        if texture_files['name'].count(texture_files['name'][number]) > 1:
            iterantTexturePath.append(texture_files['path'][number])
    
    if len(iterantTexturePath) > 0:
        _str = "发现同名贴图 : 名称如下\n"
        for Path in iterantTexturePath:
            _str += "\t%s\t%s\n" % (os.path.basename(Path),os.path.getsize(Path))
        print(_str)
        os.system('pause')
    
    mdx_files.sort()

    # 处理模型
    for filename in mdx_files:
        model = GetModelValue(filename)
        if model.type != "model":
            error_model.append((model.filename,"isNotMdxfile"))
            del model
            continue

        # 处理无贴图模型马甲类
        if model.TEXS.is_None:
            new_name = str(modelNumber).rjust(6,"0") + " " + str(time_number) + GetFileName(model.filename)[-4:]
            model.ExportModel(MapPath + '\\lop\\resource\\' + new_name)
            change_Table[2].append(new_name)
            change_Table[3].append(model.filename)
            modelNumber += 1
            continue

        modelpath = GetFilePath(model.filename)
        textures = model.TEXS.GetTexturePath()
        dirs_paths = is_sam_path(modelpath)

        # 处理错误模型
        r_tuple = is_error_model(textures,dirs_paths)
        if r_tuple and r_tuple[0]:
            error_model.append((model.filename,r_tuple[1]))
            del model
            continue

        # 处理需要变更的贴图(剔除原生贴图) 更改模型的贴图路径 变更贴图名 (空白名不需要变更 原生贴图不需要变更) 移动到主目录
        ViolenceReviseTexturePath(model,dirs_paths,textures)
        new_name = str(modelNumber).rjust(6,"0") + " " + str(time_number) + GetFileName(model.filename)[-4:]
        model.ExportModel(MapPath + '\\lop\\resource\\' + new_name)
        change_Table[2].append(new_name)
        change_Table[3].append(model.filename)
        modelNumber += 1
        print(str(len(mdx_files)) + "/" + str(mdx_files.index(filename)+1))
        if modelNumber >= groupsize and is_ModelGroup or len(mdx_files) == mdx_files.index(filename) + 1:
            UnpackMap(MapPath,MapName,is_ModelGroup,ModelGroupNumber)
            modelNumber = 1
            textureNumber = 0
            ModelGroupNumber += 1
            time_number = int(time())
            change_Table = ([],[],[],[])
            error_model = []
            os.chdir(work_path)


groupsize = 1200
skiplist = []
r_arg1 = 'G:\\素材汇集\\模型地图\\单位模型\\约战_FGO_等日漫风\\动漫类型'
r_argz = 'G:\\素材汇集\\制作模型图打包点'

main(r_arg1,r_argz,True)

# 可以通过设置 ViolenceReviseTexturePath 来选择是强行打包还是复制文件
